import serial
import yaml
import sys
import threading
import time
import struct
from scservo_sdk import * 
import numpy as np


def _float2bytes(f):
    return struct.pack('>f', f)  # '>f' 表示大端序 32 位浮点数

def _bytes2float(b):
    return struct.unpack('>f', b)[0]

def _int2bytes(i):
    return struct.pack('>H', i)  # '>H' 表示大端序 16 位无符号整数

def _bytes2int(b):
    return struct.unpack('>H', b)[0]

def _int2byte(i):
    return struct.pack('B', i)  # 'B' 表示无符号字节

def _byte2int(b):
    return struct.unpack('B', b)[0]


class _OverwritingCircularQueue:
    def __init__(self, capacity=1024):
        """初始化队列，设置容量"""
        self.capacity = capacity
        self.queue = [None] * capacity
        self.head = 0  # 指向队列头部（最旧元素）
        self.tail = 0  # 指向下一个可插入位置
        self.size = 0  # 当前队列中的元素数量
        self.lock = threading.Lock()  # 线程锁

    def enqueue(self, item):
        """非阻塞入队操作，支持覆盖"""
        with self.lock:
            self.queue[self.tail] = item
            self.tail = (self.tail + 1) % self.capacity
            if self.size == self.capacity:  # 队列已满，覆盖最旧元素
                self.head = (self.head + 1) % self.capacity
            else:
                self.size += 1
                
    def enqueue_list(self, items):
        """非阻塞批量入队操作，支持覆盖"""
        with self.lock:
            for item in items:
                self.queue[self.tail] = item
                self.tail = (self.tail + 1) % self.capacity
                if self.size == self.capacity:  # 队列已满，覆盖最旧元素
                    self.head = (self.head + 1) % self.capacity
                else:
                    self.size += 1


    def dequeue(self):
        """非阻塞出队操作"""
        with self.lock:
            if self.is_empty():
                return False, None  # 队列为空，出队失败
            item = self.queue[self.head]
            self.queue[self.head] = None
            self.head = (self.head + 1) % self.capacity
            self.size -= 1
            return True, item  # 出队成功
        
    def peek(self, index=0):
        """查询指定下标的数据"""
        with self.lock:
            if index < 0 or index >= self.size:
                return False, None  # 下标无效
            actual_index = (self.head + index) % self.capacity
            return True, self.queue[actual_index]  # 返回指定下标的数据

    def is_empty(self):
        """判断队列是否为空"""
        return self.size == 0

    def is_full(self):
        """判断队列是否已满"""
        return self.size == self.capacity

    def get_size(self):
        """获取队列中元素数量"""
        return self.size

    def __repr__(self):
        """打印队列内容，便于调试"""
        with self.lock:
            return f"Queue: {self.queue}, Head: {self.head}, Tail: {self.tail}, Size: {self.size}"
        

class AsyncTeacherConnection:
    FRAME_HEADER = 0xAA
    FRAME_TAIL   = 0xBB
    
    def __init__(self, yml_path):
        try:
            with open(yml_path, 'r') as f:
                config = yaml.safe_load(f)
        except Exception as e:
            print('Failed to load yml file: {}'.format(e))
            sys.exit(1)
            return
        
        if 'teacher' not in config:
            print('Config file must have a "teacher" section')
            sys.exit(1)
            return
        
        cfg = config['teacher']
        if 'port' not in cfg:
            print('Config file must have a "port" key in "teacher" section')
            sys.exit(1)
            return
        
        self.__port = cfg['port']    
        self.__baudrate = cfg.get('baudrate', 1500000)
        try:
            self.__serial = serial.Serial(self.__port, self.__baudrate)
        except Exception as e:
            print('Failed to open serial port: {}'.format(e))
            sys.exit(1)
            return
        
        self.__rx_queue = _OverwritingCircularQueue(256)
        self.__cb_queue = _OverwritingCircularQueue(2)
        
        self.__rx_thread = threading.Thread(target=self.__rx_thread_func, daemon=True)
        self.__rx_thread.start()
        
        self.__parse_thread = threading.Thread(target=self.__parse_thread_func, daemon=True)
        self.__parse_thread.start()
        
        self.__joint_thread = threading.Thread(target=self.__joint_thread_func, daemon=True)
        self.__joint_thread.start()
        
        self.__subcribers = {}
        self.__time = {}
        
    def __rx_thread_func(self):
        while self.__serial.is_open:
            try:
                if self.__serial.in_waiting == 0:
                    # time.sleep(0.001)
                    continue
                data = self.__serial.read(self.__serial.in_waiting)
                
                if len(data) == 0:
                    continue
                self.__rx_queue.enqueue_list(list(data))
            except Exception as e:
                print('Serial read error: {}'.format(e))
                continue
    
    def __parse_thread_func(self):
        while self.__serial.is_open:
            if self.__rx_queue.get_size() < 0:
                time.sleep(0.01)
                continue
            # header
            if self.__rx_queue.peek(0)[1] != self.FRAME_HEADER or self.__rx_queue.peek(1)[1] != self.FRAME_HEADER:
                self.__rx_queue.dequeue()
                continue
            # cmd
            cmd = self.__rx_queue.peek(2)[1]
            # data length
            data_len = self.__rx_queue.peek(3)[1]
            # check data length
            if self.__rx_queue.get_size() < data_len + 6:
                time.sleep(0.01)
                continue
            # check tail
            if self.__rx_queue.peek(data_len + 5)[1] != self.FRAME_TAIL:
                self.__rx_queue.dequeue()
                continue
            # check sum
            sum = 0
            sum += cmd
            sum += data_len
            for i in range(data_len):
                sum += self.__rx_queue.peek(4 + i)[1]
            sum = sum & 0xFF
            
            if sum != self.__rx_queue.peek(4 + data_len)[1]:
                self.__rx_queue.dequeue()
                continue

            # parse success
            tmp = []
            for i in range(data_len + 6):
                tmp.append(self.__rx_queue.dequeue()[1])
            
            self.__do_work(tmp)
    
    def __do_work(self, data):
        cmd = data[2]
        # print('data: {}'.format(data))
        if cmd == 0x01:
            # if self.__subcribers.get('joint_angles') is None:
            #     return
            
            len = data[3]
            angles = []
            for i in range(int(len / 4)):
                angle = _bytes2float(bytes(data[(4 + 4 * i):(4 + 4 * i + 4)]))
                angles.append(angle)                
            self.__cb_queue.enqueue(angles)
            
            # for callback in self.__subcribers['joint_angles']:
            #     callback(angles)
            
    def __joint_thread_func(self):
        while self.__serial.is_open:
            if self.__cb_queue.is_empty():
                time.sleep(0.01)
                continue
            angles = self.__cb_queue.dequeue()[1]
            if self.__subcribers.get('joint_angles') is None:
                continue
            for callback in self.__subcribers['joint_angles']:
                callback(angles)
                
    def __gen_frame(self, cmd, data: list):
        frame = bytearray()
        frame.append(self.FRAME_HEADER)
        frame.append(self.FRAME_HEADER)
        frame.append(cmd)
        frame.append(len(data))
        for d in data:
            frame.append(d)
        sum = 0
        sum += cmd
        sum += len(data)
        for d in data:
            sum += d
        sum = sum & 0xFF
        frame.append(sum)
        frame.append(self.FRAME_TAIL)
        return frame
        
    def request_joint_angles(self, callback):
        # if self.__time.get('joint_angles') is None:
        #     self.__time['joint_angles'] = 0
            
        # if time.time() - self.__time['joint_angles'] < 0.1:
        #     return
        # self.__time['joint_angles'] = time.time()
        
        frame = self.__gen_frame(0x01, [])
        
        if callback is not None:
            if self.__subcribers.get('joint_angles') is None:
                self.__subcribers['joint_angles'] = []
                
            if callback not in self.__subcribers['joint_angles']:
                self.__subcribers['joint_angles'].append(callback)
        
        self.__serial.write(frame)
        

class SyncTeacherConnection:
    FRAME_HEADER = 0xAA
    FRAME_TAIL   = 0xBB
    
    def __init__(self, yml_path):
        try:
            with open(yml_path, 'r') as f:
                config = yaml.safe_load(f)
        except Exception as e:
            print('Failed to load yml file: {}'.format(e))
            sys.exit(1)
            return
        
        if 'teacher' not in config:
            print('Config file must have a "teacher" section')
            sys.exit(1)
            return
        
        cfg = config['teacher']
        if 'port' not in cfg:
            print('Config file must have a "port" key in "teacher" section')
            sys.exit(1)
            return
        
        self.__port = cfg['port']    
        self.__baudrate = cfg.get('baudrate', 1000000)
        try:
            self.__serial = serial.Serial(self.__port, self.__baudrate, timeout=0.1)
        except Exception as e:
            print('Failed to open serial port: {}'.format(e))
            sys.exit(1)
            return
        
    def __gen_frame(self, cmd, data: list):
        frame = bytearray()
        frame.append(self.FRAME_HEADER)
        frame.append(self.FRAME_HEADER)
        frame.append(cmd)
        frame.append(len(data))
        for d in data:
            frame.append(d)
        sum = 0
        sum += cmd
        sum += len(data)
        for d in data:
            sum += d
        sum = sum & 0xFF
        frame.append(sum)
        frame.append(self.FRAME_TAIL)
        return frame
            
    def request_joint_angles(self):
        frame = self.__gen_frame(0x01, [])
        self.__serial.reset_input_buffer()
        self.__serial.write(frame)
        
        data = self.__serial.read(30)
        if data is None:
            return None
        arr = list(data)
        if len(arr) == 0:
            return None
        if arr[0] != self.FRAME_HEADER or arr[1] != self.FRAME_HEADER or arr[2] != 0x01 or arr[29] != self.FRAME_TAIL:
            return None
        sum = 0
        for i in range(2, 28):
            sum += arr[i]
        sum = sum & 0xFF
        if sum != arr[28]:
            return None
        
        angles = []
        for i in range(int(arr[3] / 4)):
            angle = _bytes2float(bytes(data[(4 + 4 * i):(4 + 4 * i + 4)]))
            angles.append(angle)
        return angles
        
        
class RealMachineConnection:
    ADDR_SCS_PRESENT_POSITION = 56  # 舵机当前位置的地址
    ADDR_SCS_GOAL_POSITION = 42     # 舵机目标位置的地址
    ADDR_SCS_TORQUE_ENABLE = 40     # 舵机扭矩使能地址
    
    def __init__(self, yml_path):
        try:
            with open(yml_path, 'r') as f:
                config = yaml.safe_load(f)
        except Exception as e:
            print('Failed to load yml file: {}'.format(e))
            sys.exit(1)
            return
        
        if 'real' not in config:
            print('Config file must have a "real" section')
            sys.exit(1)
            return
        
        cfg = config['real']
        if 'port' not in cfg:
            print('Config file must have a "port" key in "real" section')
            sys.exit(1)
            return
        
        self.__port = cfg['port']    
        self.__baudrate = cfg.get('baudrate', 1000000)
        
        self.__handler = PortHandler(self.__port)
        self.__handler.setBaudRate(self.__baudrate)
        
        if self.__handler.openPort() is False:
            print('Failed to open real machine port')
            sys.exit(1)
            return
        self.__scs_ids = [1, 2, 3, 4, 5, 6]
        self.__packet = PacketHandler(0)
        
    def set_joint_angles(self, angles):
        if len(angles) != 6:
            print('Invalid joint angles')
            return
        values = []
        for i in range(6):
            value = int((angles[i] + 180) * 4096 / 360)
            values.append(value)
        
        for id in self.__scs_ids:
            print('ID: {} - Setting goal position: {} {}'.format(id, angles[id - 1], SCS_LOWORD(values[id - 1])))
            result, error = self.__packet.write2ByteTxRx(self.__handler, id, self.ADDR_SCS_GOAL_POSITION, SCS_LOWORD(values[id - 1]))
            if result != COMM_SUCCESS:
                print('ID: {} - Failed to set goal position: {}'.format(id, self.__packet.getTxRxResult(result)))
            elif error != 0:
                print('ID: {} - Error code: {}'.format(id, self.__packet.getRxPacketError(error)))
            
            

class ListMachineConnection:
    ADDR_SCS_PRESENT_POSITION = 56  # 舵机当前位置的地址
    ADDR_SCS_GOAL_POSITION = 42     # 舵机目标位置的地址
    ADDR_SCS_TORQUE_ENABLE = 40     # 舵机扭矩使能地址
    
    def __init__(self, yml_path):
        try:
            with open(yml_path, 'r') as f:
                config = yaml.safe_load(f)
        except Exception as e:
            print('Failed to load yml file: {}'.format(e))
            sys.exit(1)
            return
        
        if 'list' not in config:
            print('Config file must have a "real" section')
            sys.exit(1)
            return
        
        cfg = config['list']
        if 'ports' not in cfg:
            print('Config file must have a "ports" key in "list" section')
            sys.exit(1)
            return
        
        self.__ports = cfg['ports']    
        self.__baudrate = cfg.get('baudrate', 1000000)
        
        self.__handlers = []
        self.__packets = []
        for port in self.__ports:
            handler = PortHandler(port)
            self.__handlers.append(handler)
            handler.setBaudRate(self.__baudrate)

            if handler.openPort() is False:
                print('Failed to open LIST machine port {}'.format(port))
                sys.exit(1)
                return
            
            packet = PacketHandler(0)
            self.__packets.append(packet)
        
        self.__scs_ids = [1, 2, 3, 4, 5, 6]
        
    def set_joint_angles(self, angles):
        if len(angles) != 6:
            print('Invalid joint angles')
            return
        values = []
        for i in range(6):
            value = int((angles[i] + 180) * 4096 / 360)
            values.append(value)
        
        for i in range(len(self.__handlers)):       
            handler = self.__handlers[i]     
            packet = self.__packets[i]
            
            for id in self.__scs_ids:
                # print('ID: {} - Setting goal position: {} {}'.format(id, angles[id - 1], SCS_LOWORD(values[id - 1])))
                result, error = packet.write2ByteTxRx(handler, id, self.ADDR_SCS_GOAL_POSITION, SCS_LOWORD(values[id - 1]))
                if result != COMM_SUCCESS:
                    print('ID: {} - Failed to set goal position: {}'.format(id, packet.getTxRxResult(result)))
                elif error != 0:
                    print('ID: {} - Error code: {}'.format(id, packet.getRxPacketError(error)))
            
        
if __name__ == "__main__":
    # teacher = SyncTeacherConnection('config.yml')
    # angles = teacher.request_joint_angles()
    # print(angles)
    
    # values = np.deg2rad(angles)
    values = [-30, -90, 90, 90, 0, 0]
    
    device = RealMachineConnection('config.yml')
    device.set_joint_angles(values)
         